use arrow::array::{ArrayRef, StringArray, Float64Array, Date64Array};
use arrow::datatypes::{Schema, Field, DataType as ArrowDataType};
use arrow::record_batch::RecordBatch;
use arrow::ipc::{reader, writer};
use serde_json;
use std::sync::Arc;
use std::collections::HashMap;
use std::io::{Cursor, Write};
use chrono::{DateTime, Utc};
use anyhow::Result;
use arrow::error::ArrowError;
use std::io;



pub mod types;
pub mod converters;
pub use types::*;
pub use converters::*;

/// Arrow序列化器，用于将Rust数据结构转换为Arrow格式
pub struct ArrowSerializer {
    schema: Arc<Schema>,
}

impl ArrowSerializer {
    /// 创建一个新的序列化器
    pub fn new(schema: Schema) -> Self {
        Self {
            schema: Arc::new(schema),
        }
    }

    /// 序列化记录批次
    pub fn serialize_batch(&self, data: Vec<ArrayRef>) -> Result<RecordBatch, arrow::error::ArrowError> {
        RecordBatch::try_new(self.schema.clone(), data)
    }
    
    /// 将RecordBatch转换为字节数组
    pub fn batch_to_bytes(&self, batch: &RecordBatch) -> Result<Vec<u8>, arrow::error::ArrowError> {
        let mut buf = Vec::new();
        {
            let mut writer = writer::StreamWriter::try_new(&mut buf, batch.schema().as_ref())?;
            writer.write(batch)?;
            writer.finish()?;
        }
        Ok(buf)
    }
    
    /// 将MarketData转换为RecordBatch
    pub fn market_data_to_batch(&self, data: &[MarketData]) -> Result<RecordBatch, arrow::error::ArrowError> {
        // 创建各列数据
        let timestamps: Vec<i64> = data.iter()
            .map(|d| d.timestamp.timestamp_millis())
            .collect();
        let symbols: Vec<&str> = data.iter()
            .map(|d| d.symbol.as_str())
            .collect();
        let opens: Vec<f64> = data.iter()
            .map(|d| d.open)
            .collect();
        let highs: Vec<f64> = data.iter()
            .map(|d| d.high)
            .collect();
        let lows: Vec<f64> = data.iter()
            .map(|d| d.low)
            .collect();
        let closes: Vec<f64> = data.iter()
            .map(|d| d.close)
            .collect();
        let volumes: Vec<f64> = data.iter()
            .map(|d| d.volume)
            .collect();
        let amounts: Vec<f64> = data.iter()
            .map(|d| d.amount)
            .collect();
            
        // 创建Arrow数组
        let timestamp_array = Arc::new(Date64Array::from(timestamps)) as ArrayRef;
        let symbol_array = Arc::new(StringArray::from(symbols)) as ArrayRef;
        let open_array = Arc::new(Float64Array::from(opens)) as ArrayRef;
        let high_array = Arc::new(Float64Array::from(highs)) as ArrayRef;
        let low_array = Arc::new(Float64Array::from(lows)) as ArrayRef;
        let close_array = Arc::new(Float64Array::from(closes)) as ArrayRef;
        let volume_array = Arc::new(Float64Array::from(volumes)) as ArrayRef;
        let amount_array = Arc::new(Float64Array::from(amounts)) as ArrayRef;
        
        // 组合为RecordBatch
        let arrays = vec![
            timestamp_array, 
            symbol_array, 
            open_array, 
            high_array, 
            low_array, 
            close_array, 
            volume_array, 
            amount_array
        ];
        
        // 添加元数据
        let mut metadata = HashMap::new();
        metadata.insert("data_type".to_string(), "market_data".to_string());
        
        let schema = Schema::new_with_metadata(
            vec![
                Field::new("timestamp", ArrowDataType::Date64, false),
                Field::new("symbol", ArrowDataType::Utf8, false),
                Field::new("open", ArrowDataType::Float64, false),
                Field::new("high", ArrowDataType::Float64, false),
                Field::new("low", ArrowDataType::Float64, false),
                Field::new("close", ArrowDataType::Float64, false),
                Field::new("volume", ArrowDataType::Float64, false),
                Field::new("amount", ArrowDataType::Float64, false),
            ],
            metadata
        );
        
        RecordBatch::try_new(Arc::new(schema), arrays)
    }
    
    /// 将任何可序列化的数据转换为JSON字节
    pub fn to_json_bytes<T: serde::Serialize>(&self, data: &T, data_type: DataType) -> Result<Vec<u8>, serde_json::Error> {
        let metadata = HashMap::from([
            ("data_type".to_string(), format!("{:?}", data_type).to_lowercase())
        ]);
        
        let arrow_data = ArrowData {
            data: data,
            metadata,
        };
        
        let json_str = serde_json::to_string(&arrow_data)?;
        Ok(json_str.into_bytes())
    }
}

/// Arrow反序列化器，用于将Arrow格式转换为Rust数据结构
pub struct ArrowDeserializer {
    schema: Arc<Schema>,
}

impl ArrowDeserializer {
    /// 创建一个新的反序列化器
    pub fn new(schema: Schema) -> Self {
        Self {
            schema: Arc::new(schema),
        }
    }

    /// 从RecordBatch提取列数据
    pub fn deserialize_batch(&self, batch: &RecordBatch) -> Result<Vec<ArrayRef>, arrow::error::ArrowError> {
        Ok(batch.columns().to_vec())
    }
    
    /// 从字节数组转换为RecordBatch
    pub fn bytes_to_batch(&self, bytes: &[u8]) -> Result<RecordBatch, arrow::error::ArrowError> {
        let cursor = Cursor::new(bytes);
        let mut reader = reader::StreamReader::try_new(cursor, None)?;
        
        match reader.next() {
            Some(result) => result,
            None => {
                let error = io::Error::new(io::ErrorKind::UnexpectedEof, "Empty Arrow IPC stream");
                Err(ArrowError::IoError(
                    error.to_string(),
                    error
                ))
            }
        }
    }
    
    /// 从RecordBatch转换为MarketData向量
    pub fn batch_to_market_data(&self, batch: &RecordBatch) -> Result<Vec<MarketData>, arrow::error::ArrowError> {
        let row_count = batch.num_rows();
        let mut result = Vec::with_capacity(row_count);
        
        // 获取各列数据
        let timestamp_array = batch.column(0)
            .as_any()
            .downcast_ref::<Date64Array>()
            .ok_or_else(|| arrow::error::ArrowError::ParseError(
                "Failed to downcast timestamp column".to_string()
            ))?;
            
        let symbol_array = batch.column(1)
            .as_any()
            .downcast_ref::<StringArray>()
            .ok_or_else(|| arrow::error::ArrowError::ParseError(
                "Failed to downcast symbol column".to_string()
            ))?;
            
        let open_array = batch.column(2)
            .as_any()
            .downcast_ref::<Float64Array>()
            .ok_or_else(|| arrow::error::ArrowError::ParseError(
                "Failed to downcast open column".to_string()
            ))?;
            
        let high_array = batch.column(3)
            .as_any()
            .downcast_ref::<Float64Array>()
            .ok_or_else(|| arrow::error::ArrowError::ParseError(
                "Failed to downcast high column".to_string()
            ))?;
            
        let low_array = batch.column(4)
            .as_any()
            .downcast_ref::<Float64Array>()
            .ok_or_else(|| arrow::error::ArrowError::ParseError(
                "Failed to downcast low column".to_string()
            ))?;
            
        let close_array = batch.column(5)
            .as_any()
            .downcast_ref::<Float64Array>()
            .ok_or_else(|| arrow::error::ArrowError::ParseError(
                "Failed to downcast close column".to_string()
            ))?;
            
        let volume_array = batch.column(6)
            .as_any()
            .downcast_ref::<Float64Array>()
            .ok_or_else(|| arrow::error::ArrowError::ParseError(
                "Failed to downcast volume column".to_string()
            ))?;
            
        let amount_array = batch.column(7)
            .as_any()
            .downcast_ref::<Float64Array>()
            .ok_or_else(|| arrow::error::ArrowError::ParseError(
                "Failed to downcast amount column".to_string()
            ))?;
            
        // 构建MarketData对象
        for i in 0..row_count {
            // 创建DateTime<Utc>，使用from_utc方法代替已弃用的from_timestamp_opt
            let timestamp_millis = timestamp_array.value(i);
            let secs = timestamp_millis / 1000;
            let nsecs = ((timestamp_millis % 1000) * 1_000_000) as u32;
            
            let naive_dt = chrono::NaiveDateTime::from_timestamp_opt(secs, nsecs).unwrap_or_default();
            let timestamp = DateTime::<Utc>::from_utc(naive_dt, Utc);
                
            let market_data = MarketData {
                timestamp,
                symbol: symbol_array.value(i).to_string(),
                open: open_array.value(i),
                high: high_array.value(i),
                low: low_array.value(i),
                close: close_array.value(i),
                volume: volume_array.value(i),
                amount: amount_array.value(i),
            };
            
            result.push(market_data);
        }
        
        Ok(result)
    }
    
    /// 从JSON字节转换为任何可反序列化的数据
    pub fn from_json_bytes<T: serde::de::DeserializeOwned>(&self, bytes: &[u8]) -> Result<(T, DataType), serde_json::Error> {
        let arrow_data: ArrowData<T> = serde_json::from_slice(bytes)?;
        
        // 解析数据类型
        let data_type_str = arrow_data.metadata.get("data_type")
            .cloned()
            .unwrap_or_else(|| "market_data".to_string());
            
        let data_type = match data_type_str.as_str() {
            "market_data" => DataType::MarketData,
            "technical_indicator" => DataType::TechnicalIndicator,
            "prediction_result" => DataType::PredictionResult,
            "model_params" => DataType::ModelParams,
            "error" => DataType::Error,
            _ => DataType::MarketData,
        };
        
        Ok((arrow_data.data, data_type))
    }
}